import os.path

from torchvision.datasets import ImageFolder
from torchvision.transforms import transforms
from torch.utils.data import DataLoader


def get_dataloader(root, bs, num_workers=0):
    train_transform = transforms.Compose([
        transforms.Resize(size=(250, 250)),
        transforms.RandomHorizontalFlip(),
        transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.1),
        transforms.RandomRotation(degrees=5),
        transforms.RandomResizedCrop(size=(224, 224)),
        transforms.ToTensor(),
        transforms.RandomErasing(),  # 处理Tensor，故置于ToTensor后
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    ])

    test_transform = transforms.Compose([
        transforms.Resize(size=(224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    ])
    train_dataset = ImageFolder(root=os.path.join(root, 'train'), transform=train_transform)
    test_dataset = ImageFolder(root=os.path.join(root, 'test'), transform=test_transform)

    trainloader = DataLoader(train_dataset, batch_size=bs, shuffle=True, num_workers=num_workers, pin_memory=True, prefetch_factor=2, persistent_workers=True)  # TD:改为True
    testloader = DataLoader(test_dataset, batch_size=bs, shuffle=True, num_workers=num_workers, pin_memory=True, prefetch_factor=2, persistent_workers=True)
    print(f"loader iter:\ttrian:{trainloader.__len__()}\ttest:{testloader.__len__()}")
    return trainloader, testloader


if __name__ == "__main__":
    train_loader, test_loader = get_dataloader(root=r'D:\WorkSpace\Animals10\dataset', bs=8)

    data = next(iter(test_loader))
    from visual.tensor_visual import tensor_visual

    tensor_visual(data[0], nrow=4)
